package com.totoro.mybatis.sharding.plugin;


import com.google.code.shardbatis.strategy.ShardStrategy;
import com.google.common.collect.Maps;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParser;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.dom4j.Document;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;
import org.springframework.util.CollectionUtils;

import javax.xml.parsers.SAXParser;
import javax.xml.parsers.SAXParserFactory;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringReader;
import java.util.List;
import java.util.Map;
import java.util.Properties;

/**
 * 分表插件
 * Created by Niki on 2018/8/22 19:28
 */
@Slf4j
@Intercepts({@org.apache.ibatis.plugin.Signature(type = StatementHandler.class, method = "prepare", args = {java.sql.Connection.class, java.lang.Integer.class})})
public class ShardingPlugin implements Interceptor {

    /**
     * 表库对应的容器
     */
    @Setter @Getter
    private final Map<String, String> table2DB = Maps.newHashMap();

    /**
     * 表的分表分库策略容器
     */
    @Setter
    private final Map<String, ShardStrategy> strategies = Maps.newHashMap();

    /**
     * 执行sql前的拦截，处理分库分表策略
     * * @param invocation
     *
     * @return
     * @throws Throwable
     */
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        try {
            if (CollectionUtils.isEmpty(table2DB)) {
                throw new IllegalArgumentException("请初始化逻辑表 与数据库映射");
            }
            BoundSql boundSql = getBoundSql(invocation);
            // 进行分表分库，重新组装sql
            List<String> tables = parserTableName(boundSql);


            return invocation.proceed();
        } finally {

        }

    }


    @Override
    public Object plugin(Object o) {
        return Plugin.wrap(o, this);
    }

    @Override
    public void setProperties(Properties properties) {
        String configsLocation = properties.getProperty("configsLocation");
        if (configsLocation == null) {
            throw new IllegalArgumentException("ShardPlugin[" + getClass().getName() + "] Property[configsLocation] Cannot Empty");
        }
        ClassLoader classLoader = this.getClass().getClassLoader();

        InputStream configInputStream = null;
        InputStream validateInputStream = null;
        InputStream xsdInputStream = null;
        try {

            String clazzName = this.getClass().getName();
            String xsdPath = clazzName.substring(0, clazzName.lastIndexOf('.') + 1).replace('.', '/') + "mybatis-sharding-config.xsd";

            xsdInputStream = classLoader.getResourceAsStream(xsdPath);
            configInputStream = classLoader.getResourceAsStream(configsLocation);
            validateInputStream = classLoader.getResourceAsStream(configsLocation);
            SAXParserFactory factory = SAXParserFactory.newInstance();
            SAXParser parser = factory.newSAXParser();
            SAXReader reader = new SAXReader(parser.getXMLReader());
            Document document = reader.read(configInputStream);
            Element root = document.getRootElement();
            parseStrategies(root);
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            try {
                if (validateInputStream != null) {
                    validateInputStream.close();
                }
            } catch (IOException e) {
                // ignore
            }
            try {
                if (xsdInputStream != null) {
                    xsdInputStream.close();
                }
            } catch (IOException e) {
                // ignore
            }

            try {
                if (configInputStream != null) {
                    configInputStream.close();
                }
            } catch (IOException e) {
                // ignore
            }
        }
    }

    /**
     * 解析策略文件
     *
     * @param root 根
     */
    private void parseStrategies(Element root) throws ClassNotFoundException, InstantiationException, IllegalAccessException {
        List<?> strategies = root.elements("strategy");
        if (strategies != null) {
            for (Object o : strategies) {
                Element strategy = (Element) o;
                String logicTable = strategy.attribute("logicTable").getStringValue();
                String strategyClass = strategy.attribute("class").getStringValue();
                Class<?> clazz = Class.forName(strategyClass);
                ShardStrategy shardStrategy = (ShardStrategy) clazz.newInstance();
                if (this.strategies.containsKey(logicTable)) {
                    throw new IllegalArgumentException("LogicTable[" + logicTable + "] Duplicate");
                }
                this.strategies.put(logicTable, shardStrategy);
            }
        }
    }

    /**
     * 从拦截的StatementHandler对象中取出SQL语句对象
     * @param invocation
     * @return
     */
    private BoundSql getBoundSql(Invocation invocation) {
        // 因为拦截的是StatementHandler的对象。所以invocation是StatementHandler的实例
        StatementHandler statementHandler = (StatementHandler) (invocation.getTarget());
        return statementHandler.getBoundSql();
    }

    /**
     * 解析sql的标明
     * @param boundSql
     * @return
     * @throws JSQLParserException
     */
    private List<String> parserTableName(BoundSql boundSql) throws JSQLParserException {
        String sql = boundSql.getSql();
        log.info("拦截前的sql={}", sql);

        Select stmt = (Select) CCJSqlParserUtil.parse(sql);

        TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
        List<String> tables = tablesNamesFinder.getTableList(stmt);
        log.debug("插件拦截之前的表名={}",tables);
        return tables;
    }

}
